import os
import unittest

import torch
from torch import hub
from torch.backends import cudnn
from torchsummary import summary
from config import Config

from training.models.convnext_train import convnext_train
# from training.models.rgb_lace_convnext import rgb_lace_convnext


if torch.cuda.is_available() == False:
    # CPU:
    CONFIG = Config()
    hub.set_dir(CONFIG['WINDOWS_TORCH_HOME'])
else:
    # GPU:
    CONFIG = Config()
    hub.set_dir(CONFIG['TORCH_HOME'])
    os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG['CUDA_VISIBLE_DEVICES']
    torch.backends.cudnn.benchmark = True

class ConvnextTrainTextCase(unittest.TestCase):

    def test_summary_convnext_train(self):
        # self.assertTrue(torch.cuda.is_available())
        # model = rgb_lace_convnext(pretrained=True)
        model = convnext_train(pretrained=True)

        if torch.cuda.is_available() == False:
            model = model.cpu()
        else:
            model = model.cuda()


        # input_size = [model.default_cfg['input_size'], model.default_cfg['input_size']]
        input_size = model.default_cfg['input_size']
        summary(model, input_size=input_size)


if __name__ == '__main__':
    unittest.main()